// -----------------------------
// Author: 		Dhaenens Gert	
// Creation: 	12/19/2014 7:11:37 PM				
// -----------------------------

#include "PostPass.h"

#include "Application.h"
#include "FileManager.h"
#include "RenderManager.h"
#include "Window.h"

#include "Scene.h"
#include "Camera.h"

#define DISPATCH_NUM_THREADS_X	16
#define DISPATCH_NUM_THREADS_Y	16

struct PostBuffer
{
	Vector3 WaterOverlayColor;
	Float	WaterOverlayStrength;
	UInt32	WaterOverlayEnabled;
	Float	WaterVignetteRadius;
	UInt32	ScreenWidth;
	UInt32	ScreenHeight;
	Vector3	CameraPosition;
	Float	FogNear;
	Vector3	FogColor;
	Float	FogFar;
	Vector3 FogSecondColor;
	UInt32	FogEnabled;
};

PostPass::PostPass(RenderManager* lpManager, const RenderTargets& targets)
	: RenderPass(lpManager, "Post", targets)
	, m_lpComputeShader(nullptr)
	, m_lpPostBuffer(nullptr)
	, m_ThreadCountX(0)
	, m_ThreadCountY(0)
	, m_ThreadCountZ(0)
	, m_WaterOverlayColor(Color::Blue)
	, m_bWaterOverlayEnabled(false)
	, m_WaterOverlayStrength(0.0f)
{
}

PostPass::~PostPass()
{
	NullRelease(&m_lpPostBuffer);
	NullRelease(&m_lpComputeShader);
}

void PostPass::Init(Application* lpApp)
{
	TString computeShader = lpApp->GetFiles()->GetFileName(_T("Shader_PostCompute"));

	m_lpComputeShader = ShaderCompiler::LoadComputeFromFile(m_lpRenderManager->GetDevice(), computeShader.c_str());

	m_lpPostBuffer = IBuffer::Create(m_lpRenderManager->GetDevice(), BufferType::Constant, Usage::Default, 0, sizeof(PostBuffer), 1, nullptr);

	m_ThreadCountX = (lpApp->GetWindow()->GetWidth() + DISPATCH_NUM_THREADS_X) / DISPATCH_NUM_THREADS_X;
	m_ThreadCountY = (lpApp->GetWindow()->GetHeight() + DISPATCH_NUM_THREADS_Y) / DISPATCH_NUM_THREADS_Y;
	m_ThreadCountZ = 1;
}
void PostPass::Apply(Application* lpApp, Scene* lpScene)
{
	IRenderDevice* lpDevice = m_lpRenderManager->GetDevice();

	// Update the Post Buffer
	PostBuffer post = {};
	post.WaterOverlayColor = m_WaterOverlayColor.ToVector3();
	post.WaterOverlayStrength = m_WaterOverlayStrength;
	post.WaterOverlayEnabled = (m_bWaterOverlayEnabled) ? 1 : 0;
	post.WaterVignetteRadius = m_WaterVignetteRadius;
	post.ScreenWidth = lpApp->GetWindow()->GetWidth();
	post.ScreenHeight = lpApp->GetWindow()->GetHeight();
	post.CameraPosition = lpScene->GetCamera()->GetPosition();
	post.FogNear = m_FogNear;
	post.FogFar = m_FogFar;
	post.FogColor = m_FogColor.ToVector3();
	post.FogSecondColor = m_FogSecondColor.ToVector3();
	post.FogEnabled = (m_FogEnabled) ? 1 : 0;
	m_lpPostBuffer->Update(&post, sizeof(PostBuffer));

	lpDevice->SetRenderTarget(nullptr, nullptr);

	lpDevice->SetComputeShader(m_lpComputeShader);
	m_lpComputeShader->SetTargetAsResource(0, m_Targets.RawRenderTarget3D);
	m_lpComputeShader->SetTargetAsResource(1, m_Targets.PositionTarget);
	m_lpComputeShader->SetTargetAsResource(2, m_Targets.NormalTarget);
	m_lpComputeShader->SetTargetAsResource(3, m_Targets.AlbedoTarget);
	m_lpComputeShader->SetTargetAsResource(4, m_Targets.AmbientTarget);
	m_lpComputeShader->SetStencilAsResource(5, m_Targets.DepthStencil);
	m_lpComputeShader->SetBufferAsConstant(0, m_lpPostBuffer);
	#ifdef _NOD2D
	m_lpComputeShader->SetTargetAsUav(0, m_Targets.OutputTarget);
	#else
	m_lpComputeShader->SetTargetAsUav(0, m_Targets.PostRenderTarget3D);
	#endif

	// Dispatch the Compute Shader
	lpDevice->Dispatch(m_ThreadCountX, m_ThreadCountY, m_ThreadCountZ);

	m_lpComputeShader->ClearResource(0);
	m_lpComputeShader->ClearResource(1);
	m_lpComputeShader->ClearResource(2);
	m_lpComputeShader->ClearResource(3);
	m_lpComputeShader->ClearResource(4);
	m_lpComputeShader->ClearResource(5);
	m_lpComputeShader->ClearUav(0);
}
void PostPass::UpdateTargets(Application* lpApp, const RenderTargets& targets)
{
	RenderPass::UpdateTargets(lpApp, targets);

	// Update the Thread Counts
	m_ThreadCountX = (lpApp->GetWindow()->GetWidth() + DISPATCH_NUM_THREADS_X) / DISPATCH_NUM_THREADS_X;
	m_ThreadCountY = (lpApp->GetWindow()->GetHeight() + DISPATCH_NUM_THREADS_Y) / DISPATCH_NUM_THREADS_Y;
	m_ThreadCountZ = 1;
}

void PostPass::SetWaterOverlayEnabled(Bool bEnable)
{
	m_bWaterOverlayEnabled = bEnable;
}
void PostPass::SetWaterOverlayColor(Color color)
{
	m_WaterOverlayColor = color;
}
void PostPass::SetWaterOverlayStrength(Float strength)
{
	m_WaterOverlayStrength = strength;
}
void PostPass::SetWaterVignetteRadius(Float radius)
{
	m_WaterVignetteRadius = radius;
}
void PostPass::SetFogNear(Float fogNear)
{
	m_FogNear = fogNear;
}
void PostPass::SetFogFar(Float fogFar)
{
	m_FogFar = fogFar;
}
void PostPass::SetFogColor(const Color& color)
{
	m_FogColor = color;
}
void PostPass::SetFogSecondColor(const Color& color)
{
	m_FogSecondColor = color;
}
void PostPass::SetFogEnabled(Bool bEnabled)
{
	m_FogEnabled = bEnabled;
}